package com.poetry.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.poetry.domain.po.ChatMessage;
import com.poetry.exception.R;
import com.poetry.exception.UnauthorizedException;
import com.poetry.mapper.ChatMessageMapper;
import com.poetry.service.IChatMessageService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

@Service
public class ChatMessageServiceImpl extends ServiceImpl<ChatMessageMapper, ChatMessage>
        implements IChatMessageService, ChatMemory {

    @Override
    public void add(String conversationId, Message message) {
        IChatMessageService.super.add(conversationId, message);
    }

    @Override
    @Transactional
    public void add(String conversationId, List<Message> messages) {
        List<ChatMessage> chatMessages = messages.stream()
                .map(msg -> {
                    ChatMessage chatMessage = new ChatMessage();
                    chatMessage.setConversationId(conversationId);
                    chatMessage.setContent(msg.getText());
                    if (msg instanceof UserMessage) {
                        chatMessage.setType(ChatMessage.MessageType.USER);
                    }
                    if (msg instanceof AssistantMessage) {
                        chatMessage.setType(ChatMessage.MessageType.ASSISTANT);
                    }
                    chatMessage.setCreateTime(LocalDateTime.now());
                    return chatMessage;
                }).toList();
        this.saveBatch(chatMessages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<ChatMessage> chatMessages = this.query().eq("conversation_id", conversationId)
                .orderByDesc("create_time")
                .last("limit " + lastN).list();
        Collections.reverse(chatMessages);
        return chatMessages.stream()
                .map(chatMessage -> {
                    switch (chatMessage.getType()) {
                        case USER:
                            return new UserMessage(chatMessage.getContent());
                        case ASSISTANT:
                            return new AssistantMessage(chatMessage.getContent());
                        default:
                            throw new UnauthorizedException("未知消息类型");
                    }
                }).collect(Collectors.toList());
    }

    @Override
    @Transactional
    public void clear(String conversationId) {
        this.remove(new LambdaQueryWrapper<ChatMessage>()
                .eq(ChatMessage::getConversationId, conversationId));
    }

    @Override
    public R<List<ChatMessage>> getChatHistory(String conversationId) {
        List<ChatMessage> chatHistory = this.query()
                .eq("conversation_id", conversationId)
                .last("limit 10")
                .list();
        if (chatHistory == null || chatHistory.isEmpty()) return R.success("对话信息为空", List.of());
        return R.success("对话记录获取成功", chatHistory);
    }
}
